import joblib
import numpy as np
import gradio as gr
from PIL import Image
from skimage.transform import resize

# 加载模型和标准化器
model_path = "C:/Users/86178/Desktop/faiss_dog_cat_question-main/best_model/best_model.joblib"
scaler_path = "C:/Users/86178/Desktop/faiss_dog_cat_question-main/best_model/scaler.joblib"

model = joblib.load(model_path)
scaler = joblib.load(scaler_path)

# 检查标准化器的特征数量
print(f"StandardScaler expects {scaler.n_features_in_} features")

# 图像预处理函数
def preprocess_image(image):
    image = np.array(image.convert('L'))  # 将图像转换为灰度图像
    image = resize(image, (32, 32), anti_aliasing=True)  # 调整图像大小至 32x32
    image = image.flatten().astype('float32')  # 展平图像并转换为浮点型
    image = scaler.transform([image])  # 使用scaler进行标准化
    return image

# 推理函数
def predict(image):
    processed_image = preprocess_image(image)
    prediction = model.predict(processed_image)
    probability = model.predict_proba(processed_image)[0]

    if prediction[0] == 0:
        label = "Cat"
    else:
        label = "Dog"

    return label, str(probability)  # 返回两个独立的值

# 创建 Gradio 界面
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Textbox(label="Prediction"), gr.Textbox(label="Probability")]
)

# 启动应用
iface.launch()